import os
import cv2
from torch.utils.data import DataLoader
import argparse
import glob
from dudenet import *
import os.path
from torch.autograd import Variable
from utils import *
import torch

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

parser = argparse.ArgumentParser(description="Net2_Test")
parser.add_argument("--logdir", type=str, default="sigma_50", help='path of log files')
# parser.add_argument("--test_data", type=str, default='Set68_cropped', help='test on Set12 or Set68')
parser.add_argument("--test_data", type=str, default='Set12', help='test on Set12 or Set68')
parser.add_argument("--test_noiseL", type=float, default=50, help='noise level used on test set')
opt = parser.parse_args()

def normalize(data):
    return data/255.

def main():
    # Build model
    # torch.seed()
    save_dir = 'sigma' + str(opt.test_noiseL) + '_' + opt.test_data + '/'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    print('Loading model ...\n')
    net = dudenet()
    device_ids = [0]
    model = nn.DataParallel(net, device_ids=device_ids).cuda()
    model.load_state_dict(torch.load(os.path.join(opt.logdir, 'epoch2_13.pth')))
    model.eval()
    # load data info
    print('Loading data info ...\n')
    files_source = glob.glob(os.path.join('./data/test/', opt.test_data, '*.*'))
    files_source.sort()
    # process data
    psnr_test, ssim_test = 0, 0
    for f in files_source:
        # image
        Img = cv2.imread(f)
        f2 = replace
        ref_Img = cv2.imread(f2)
        Img = normalize(np.float32(Img[:, :, 0]))
        Img = np.expand_dims(Img, 0)
        Img = np.expand_dims(Img, 1)
        ISource = torch.Tensor(Img)
        # noise
        # torch.manual_seed(0) #set the seed
        noise = torch.FloatTensor(ISource.size()).normal_(mean=0, std=opt.test_noiseL/255.)
        # noisy image
        INoisy = torch.clamp(ISource + noise, 0., 1.).cuda()
        # ISource, INoisy = Variable(ISource.cuda()), Variable(INoisy.cuda())
        with torch.no_grad():  # this can save much memory
            Out = torch.clamp(model(INoisy), 0., 1.)
            out1 = Out[0, 0, :, :].cpu().numpy()
        psnr = batch_PSNR(Out, ISource, 1.)
        ssim =batch_SSIM(Out, ISource, 1.)
        psnr_test += psnr
        ssim_test += ssim
        print("%s PSNR %f, SSIM %f" % (f, psnr, ssim))
        a = os.path.basename(f)
        cv2.imwrite(os.path.join(save_dir, a), np.round(out1*255.0))
    psnr_test /= len(files_source)
    ssim_test /= len(files_source)
    print("\nPSNR on test data %f,SSIM on test data %f" % (psnr_test, ssim_test))


if __name__ == "__main__":
    main()


